#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
PDF转图片并按比例裁剪工具

此脚本将PDF文件转换为图片，并按照指定的比例进行裁剪。
支持四个参数分别控制左、右、上、下的裁剪比例。
默认左右保留75%，上下保留80%区域。
原始图片和裁剪后的图片都会保存在随机命名的文件夹中。
"""

import os
import argparse
import uuid
import re
from pdf2image import convert_from_path
from PIL import Image
import logging

# 配置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

def parse_args():
    """解析命令行参数"""
    parser = argparse.ArgumentParser(description='Convert PDF to images and crop by ratio')
    parser.add_argument('pdf_path', help='PDF file path')
    parser.add_argument('-o', '--output', default='output', help='Output directory, default "output"')
    parser.add_argument('-f', '--format', default='jpg', choices=['jpg', 'png'], help='Output image format, default jpg')
    parser.add_argument('-d', '--dpi', type=int, default=200, help='Conversion DPI, default 200')
    parser.add_argument('--left-ratio', type=float, default=0.75, help='Left crop ratio')
    parser.add_argument('--right-ratio', type=float, default=0.8, help='Right crop ratio')
    parser.add_argument('--top-ratio', type=float, default=0.7, help='Top crop ratio')
    parser.add_argument('--bottom-ratio', type=float, default=1.0, help='Bottom crop ratio')
    parser.add_argument('--no-crop', action='store_true', help='No cropping, convert PDF to images only')
    parser.add_argument('--merg-corp-pdf', type=str, default="1", help='Merge cropped images to PDF')

    
    return parser.parse_args()

def convert_pdf_to_images(pdf_path, output_dir, dpi=200, fmt='jpg'):
    """将PDF转换为图片"""
    logger.info(f"正在将PDF转换为图片: {pdf_path}")
    
    # 确保输出目录存在
    os.makedirs(output_dir, exist_ok=True)
    
    # 转换PDF为图片
    images = convert_from_path(pdf_path, dpi=dpi)
    
    image_paths = []
    for i, image in enumerate(images):
        # 构建输出文件路径
        output_path = os.path.join(output_dir, f"page_{i+1}.{fmt}")
        # 保存图片
        if fmt.lower() == 'jpg':
            image.save(output_path, 'JPEG')
        else:
            image.save(output_path, fmt.upper())
        image_paths.append(output_path)
        logger.info(f"已保存第 {i+1} 页: {output_path}")
    
    return image_paths

def crop_image_by_ratio(image_path, output_path, left_ratio=0.75, right_ratio=0.75, top_ratio=0.8, bottom_ratio=0.8):
    """按比例裁剪图片，使用四个参数分别控制左、右、上、下的裁剪比例"""
    logger.info(f"正在裁剪图片: {image_path}")
    
    # 打开图片
    img = Image.open(image_path)
    width, height = img.size
    
    # 计算裁剪区域
    # 左右方向：从中心开始，向左和右各扩展指定比例
    horizontal_crop_width = int(width * (left_ratio + right_ratio) / 2)
    # 上下方向：从中心开始，向上和下各扩展指定比例
    vertical_crop_height = int(height * (top_ratio + bottom_ratio) / 2)
    
    # 计算裁剪区域的坐标
    left = max(0, (width - horizontal_crop_width) // 2)
    top = max(0, (height - vertical_crop_height) // 2)
    right = min(width, left + horizontal_crop_width)
    bottom = min(height, top + vertical_crop_height)
    
    # 裁剪图片
    cropped_img = img.crop((left, top, right, bottom))
    
    # 保存裁剪后的图片
    cropped_img.save(output_path)
    logger.info(f"已保存裁剪后的图片: {output_path}")
    
    return output_path

def merge_images_to_pdf(image_paths, output_pdf_path):
    """将图片列表合成为PDF文件"""
    logger.info(f"正在将图片合成为PDF: {output_pdf_path}")
    
    if not image_paths:
        logger.warning("没有图片可以合成PDF")
        return None
    
    # 按照文件名中的数字顺序排列（page_1, page_2, page_3...）
    def natural_sort_key(path):
        """提取文件名中的数字用于排序"""
        filename = os.path.basename(path)
        # 使用正则表达式提取数字
        match = re.search(r'page_(\d+)', filename)
        if match:
            return int(match.group(1))
        return 0
    
    # 按数字顺序排列图片路径
    sorted_image_paths = sorted(image_paths, key=natural_sort_key)
    
    # 打开第一张图片作为基础
    images = []
    for image_path in sorted_image_paths:
        img = Image.open(image_path)
        # 如果图片是RGBA模式，转换为RGB模式
        if img.mode == 'RGBA':
            img = img.convert('RGB')
        elif img.mode != 'RGB':
            img = img.convert('RGB')
        images.append(img)
        logger.info(f"正在处理第 {natural_sort_key(image_path)} 页: {os.path.basename(image_path)}")
    
    # 保存为PDF
    if images:
        images[0].save(
            output_pdf_path, 
            "PDF", 
            resolution=100.0, 
            save_all=True, 
            append_images=images[1:]
        )
        logger.info(f"已保存合成的PDF: {output_pdf_path}")
        return output_pdf_path
    
    return None

def main():
    """主函数"""
    args = parse_args()
    
    # 检查PDF文件是否存在
    if not os.path.exists(args.pdf_path):
        logger.error(f"PDF文件不存在: {args.pdf_path}")
        return
    
    # 创建随机文件夹名
    random_folder_name = str(uuid.uuid4())[:8]
    base_output_dir = os.path.join(args.output, random_folder_name)
    os.makedirs(base_output_dir, exist_ok=True)
    
    # 创建原始图片目录
    original_dir = os.path.join(base_output_dir, "original")
    os.makedirs(original_dir, exist_ok=True)
    
    # 转换PDF为图片（保存到原始图片目录）
    image_paths = convert_pdf_to_images(args.pdf_path, original_dir, args.dpi, args.format)
    
    # 如果不需要裁剪，直接返回
    if args.no_crop:
        logger.info(f"已完成PDF转图片，原始图片保存在: {original_dir}")
        return
    
    # 裁剪图片
    cropped_dir = os.path.join(base_output_dir, "cropped")
    os.makedirs(cropped_dir, exist_ok=True)
    
    cropped_image_paths = []
    for image_path in image_paths:
        # 构建裁剪后的图片路径
        filename = os.path.basename(image_path)
        cropped_path = os.path.join(cropped_dir, filename)
        
        # 使用四个参数裁剪图片
        crop_image_by_ratio(
            image_path, 
            cropped_path, 
            args.left_ratio, 
            args.right_ratio, 
            args.top_ratio, 
            args.bottom_ratio
        )
        cropped_image_paths.append(cropped_path)
    
    # 合成PDF（如果需要）
    # 非0即为合并
    merge_to_pdf = getattr(args, 'merg_corp_pdf', '1') != '0'
    if merge_to_pdf:
        # 获取原始PDF文件名（不含扩展名）
        pdf_basename = os.path.splitext(os.path.basename(args.pdf_path))[0]
        merged_pdf_path = os.path.join(base_output_dir, f"{pdf_basename}_cropped.pdf")
        merge_images_to_pdf(cropped_image_paths, merged_pdf_path)
    
    logger.info(f"处理完成！")
    logger.info(f"原始图片保存在: {original_dir}")
    logger.info(f"裁剪后的图片保存在: {cropped_dir}")
    if merge_to_pdf:
        logger.info(f"合成的PDF保存在: {os.path.join(base_output_dir, f'{os.path.splitext(os.path.basename(args.pdf_path))[0]}_cropped.pdf')}")

if __name__ == "__main__":
    main()